package grpc

import (
	"TOScheduler/src/pb"
	"TOScheduler/src/util"
	"context"
	"fmt"
	"google.golang.org/grpc"
	"strconv"
	"sync"
	"time"
)

var (
	portMap map[int32]int
	cliMap  sync.Map
)

func InitAIClient(portBase int, portDiff int) {
	portMap = make(map[int32]int)
	for _, taskType := range pb.Basic_TaskType_value {
		portFactor := int(taskType)
		if portFactor > 0 {
			portFactor = portFactor - 1
		}
		portMap[taskType] = portBase + portDiff*portFactor
	}
	WaitAndProcess()
}

func Tran(basic *pb.Basic) {
	basic.OffloadTo = basic.ArriveAt
	latency := &pb.Latency{
		Deadline:   0,
		LocalDelay: 0,
		Timestamps: nil,
	}
	log.Debugf("[gRPC AI] Tran task with length (%d): %s", len(chRes), basic.Id)
	// back nil result
	chRes <- nil
	SendGrpcTOAI(basic, latency)
}

func WaitAndProcess() {
	time.Sleep(time.Second)
	for {
		log.Debugf("[gRPC AI] blocked with length %d (1), %d (2)", len(chFlag1), len(chFlag2))
		id1 := <-chFlag1
		id2 := <-chFlag2
		id := util.Min(id1, id2)
		brief0, ok1 := briefMap.LoadAndDelete(id)
		task0, ok2 := taskMap.LoadAndDelete(id)
		if !(ok1 && ok2) {
			log.Errorf("[gRPC AI] Concurrent Failure On [%d/%d] !",
				id1, id2)
			continue
		}
		go SendGrpcPrepare(brief0, task0)
	}
}

func SendGrpcPrepare(brief0 interface{}, task0 interface{}) {
	brief := brief0.(*pb.ToTOBrief)
	task := task0.(*pb.ToTOTask)
	basicSDK := task.Basic
	latencySDK := task.Latency
	basicTD := brief.Basic
	latencyTD := brief.Latency
	log.Debug("[gRPC AI] All channel data received")
	end := util.GenTimestamp()
	for i := 0; i < len(basicSDK); i++ {
		finalBasic := &pb.Basic{
			Id:        basicSDK[i].Id,
			Device:    basicSDK[i].Device,
			TaskType:  basicTD[i].TaskType,
			ArriveAt:  basicTD[i].ArriveAt,
			OffloadTo: basicTD[i].OffloadTo,
			Snapshot:  basicSDK[0].Snapshot,
		}
		timestamps := latencySDK[i].Timestamps
		timestamps = append(timestamps, end)
		finalLatency := &pb.Latency{
			Deadline:   latencyTD[i].Deadline,
			LocalDelay: latencyTD[i].LocalDelay,
			Timestamps: timestamps,
		}
		go SendGrpcTOAI(finalBasic, finalLatency)
	}
}

func SendGrpcTOAI(basic *pb.Basic, latency *pb.Latency) {
	// set the request message
	res := pb.ToAI{
		Basic:   basic,
		Latency: latency,
	}
	offloadTo := basic.OffloadTo
	if len(offloadTo) == 0 {
		log.Warnf("[gRPC AI] Blank offloadTo by: %v", basic.Id)
		offloadTo = basic.ArriveAt
	}
	aiPort := portMap[int32(basic.TaskType)]
	clientAddress := fmt.Sprintf("%s:%s", offloadTo, strconv.Itoa(aiPort))
	cliVal, ok := cliMap.Load(clientAddress)
	if !ok {
		// build connection
		conn, err := grpc.Dial(clientAddress, grpc.WithInsecure())
		if err != nil {
			log.Errorf("[gRPC AI] error ( %s ): %v", clientAddress, err)
		}
		cliVal = pb.NewSendMsgTOAIClient(conn)
		cliMap.Store(clientAddress, cliVal)
	}
	cli := cliVal.(pb.SendMsgTOAIClient)
	r, err := cli.SendMsgTOAI(context.Background(), &res)
	if err != nil {
		log.Errorf("[gRPC AI] error: ( %s ): %v", clientAddress, err)
		return
	}
	if len(r.Res.Persons) > 0 {
		log.Debugf("[gRPC AI] recv FR message with length: %v", r.Res)
		chRes <- r.Res
	}
}
